package org.jeecg.common.util;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.Security;
import java.util.Base64;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;
import java.security.SecureRandom;
import java.util.Base64;

import org.bouncycastle.jce.provider.BouncyCastleProvider;

/**
 * SM4加密工具类 - 基于特定Base64字符串生成密钥
 */
public class SM4Utils {
    static {
        // 注册Bouncy Castle安全提供者
        if (Security.getProvider(BouncyCastleProvider.PROVIDER_NAME) == null) {
            Security.addProvider(new BouncyCastleProvider());
        }
    }

    // 算法名称
    private static final String ALGORITHM = "SM4";
    // 默认模式和填充方式
    private static final String ALGORITHM_MODE_PADDING = "SM4/ECB/PKCS5Padding";
    // 密钥长度（字节）
    private static final int KEY_SIZE = 16;

    private static final int KEY_LENGTH = 16;
    // 盐值长度
    private static final int SALT_LENGTH = 16;
    // 迭代次数（提高安全性）
    private static final int ITERATIONS = 65536;

    public static byte[] deriveKey(String password, byte[] salt) throws Exception {
        PBEKeySpec spec = new PBEKeySpec(
                password.toCharArray(),
                salt,
                ITERATIONS,
                KEY_LENGTH * 8 // 转换为位数
        );

        SecretKeyFactory factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256");
        return factory.generateSecret(spec).getEncoded();
    }

    /**
     * 生成随机盐值
     */
    public static byte[] generateSalt() {
        SecureRandom random = new SecureRandom();
        byte[] salt = new byte[SALT_LENGTH];
        random.nextBytes(salt);
        return salt;
    }

    /**
     * 从Base64字符串生成SM4密钥
     * @param base64Str Base64编码的字符串
     * @return SM4密钥字节数组
     * @throws Exception 如果解码失败或长度不符合要求
     */
    public static byte[] generateKeyFromBase64(String base64Str) throws Exception {
        // 解码Base64字符串
        byte[] decodedBytes = Base64.getDecoder().decode(base64Str);

        // 验证长度是否符合SM4密钥要求（16字节）
        if (decodedBytes.length != KEY_SIZE) {
            throw new IllegalArgumentException(
                    "Base64解码后的字节长度必须为" + KEY_SIZE + "字节，但实际为" + decodedBytes.length + "字节");
        }

        return decodedBytes;
    }

    /**
     * 使用生成的密钥进行SM4加密
     * @param plainText 明文
     * @param keyBytes SM4密钥字节数组
     * @return 加密后的Base64编码字符串
     * @throws Exception 如果加密失败
     */
    public static String encrypt(String plainText, byte[] keyBytes) throws Exception {
        // 创建密钥规范
        SecretKeySpec secretKey = new SecretKeySpec(keyBytes, ALGORITHM);

        // 获取加密器实例
        Cipher cipher = Cipher.getInstance(ALGORITHM_MODE_PADDING, "BC");
        cipher.init(Cipher.ENCRYPT_MODE, secretKey);

        // 执行加密
        byte[] encryptedBytes = cipher.doFinal(plainText.getBytes(StandardCharsets.UTF_8));

        // 返回Base64编码的密文
        return Base64.getEncoder().encodeToString(encryptedBytes);
    }

    /**
     * 使用生成的密钥进行SM4解密
     * @param cipherTextBase64 加密后的Base64编码字符串
     * @param keyBytes SM4密钥字节数组
     * @return 明文字符串
     * @throws Exception 如果解密失败
     */
    public static String decrypt(String cipherTextBase64, byte[] keyBytes) throws Exception {
        // 解码Base64密文
        byte[] cipherBytes = Base64.getDecoder().decode(cipherTextBase64);

        // 创建密钥规范
        SecretKeySpec secretKey = new SecretKeySpec(keyBytes, ALGORITHM);

        // 获取解密器实例
        Cipher cipher = Cipher.getInstance(ALGORITHM_MODE_PADDING, "BC");
        cipher.init(Cipher.DECRYPT_MODE, secretKey);

        // 执行解密
        byte[] decryptedBytes = cipher.doFinal(cipherBytes);

        // 返回明文字符串
        return new String(decryptedBytes, StandardCharsets.UTF_8);
    }

    /**
     * 生成随机SM4密钥并返回Base64编码字符串
     * @return Base64编码的SM4密钥
     * @throws Exception 如果生成失败
     */
    public static String generateRandomKeyBase64() throws Exception {
        KeyGenerator kg = KeyGenerator.getInstance(ALGORITHM, "BC");
        kg.init(128);
        byte[] keyBytes = kg.generateKey().getEncoded();
        return Base64.getEncoder().encodeToString(keyBytes);
    }

    /**
     * 验证UMNA2TLYS8FDFL6T作为SM4密钥的示例
     */
    public static void main(String[] args) {
        try {
            // 你的Base64字符串
            String base64Key = "UMNA2TLYS8FDFL6T";
            byte[] salt = generateSalt();
            byte[] key = deriveKey(base64Key, salt);
            System.out.println("生成的SM4密钥字节长度1: " + key.length + " 字节");
            System.out.println("Base64编码: " + Base64.getEncoder().encodeToString(key));



            // 生成SM4密钥
           // byte[] keyBytes = generateKeyFromBase64(base64Key);
           // System.out.println("生成的SM4密钥字节长度: " + keyBytes.length);

            // 待加密的明文
            String plainText = "Hello, SM4加密测试!";

            // 加密
            String encryptedText = encrypt(plainText, key);
            System.out.println("加密后的Base64: " + encryptedText);

            // 解密
            String decryptedText = decrypt(encryptedText, key);
            System.out.println("解密后的明文: " + decryptedText);

        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}